import gradio as gr
from login import AuthenticatedApp
from openai_tools import generate_illustration
from prompt_templates import PromptTemplates, prompt_default
from tools import logger


def create_gradio_interface():
    app = AuthenticatedApp()
    templates = PromptTemplates()

    with gr.Blocks() as demo:
        session_id = gr.State("")
        current_template = gr.State("city_scene")  # 新增：记录当前模板类型
        
        with gr.Tab("账户"):
            with gr.Group():
                username_input = gr.Textbox(label="用户名")
                password_input = gr.Textbox(label="密码", type="password")
                register_btn = gr.Button("注册")
                login_btn = gr.Button("登录")
                logout_btn = gr.Button("登出")
                account_status = gr.Textbox(label="状态")
            
            def handle_register(username, password):
                result = app.register(username, password)
                if result["success"]:
                    return result["message"], result["session_id"]
                return result["message"], ""
            
            def handle_login(username, password):
                result = app.login(username, password)
                if result["success"]:
                    return result["message"], result["session_id"]
                return result["message"], ""
            
            def handle_logout(curr_session_id):
                result = app.logout(curr_session_id)
                return result["message"]
            
            register_btn.click(
                handle_register,
                inputs=[username_input, password_input],
                outputs=[account_status, session_id]
            )
            
            login_btn.click(
                handle_login,
                inputs=[username_input, password_input],
                outputs=[account_status, session_id]
            )
            
            logout_btn.click(
                handle_logout,
                inputs=[session_id],
                outputs=[account_status]
            )
        
        # 添加深度记事插画生成选项卡
        with gr.Tab("深度记事-明信片插画生成"):
            with gr.Row():
                with gr.Column(scale=1):
                    # 左侧功能按钮区
                    with gr.Group():
                        city_scene_btn = gr.Button("城市景点介绍", variant="secondary")
                    with gr.Group():
                        knowledge_btn = gr.Button("知识学习小报", variant="secondary")
                    with gr.Group():
                        custom_template_btn = gr.Button("新增提示词词模版", variant="secondary")
                    
                     # 添加空行分隔
                    gr.Markdown("&nbsp;")
                    gr.Markdown("&nbsp;")
                    gr.Markdown("&nbsp;")
                    
                    with gr.Group():
                        generate_prompt_btn = gr.Button("生成提示词", variant="primary")
                    with gr.Group():
                        generate_image_btn = gr.Button("一键生图", variant="primary")

                
                with gr.Column(scale=1):
                    # 中间参数设置区
                    with gr.Group(visible=True) as city_scene_group:
                        gr.Markdown("### 城市景点介绍")
                        city_input = gr.Textbox(label="城市", value="杭州")
                        scene_input = gr.Textbox(label="景区", value="西湖")
                    
                    with gr.Group(visible=False) as knowledge_group:
                        gr.Markdown("### 知识学习小报")
                        topic_input = gr.Textbox(label="主题", value="AI强化学习")
                    
                    with gr.Group(visible=False) as custom_template_group:
                        gr.Markdown("### 自定义提示词")
                        custom_prompt = gr.Textbox(label="自定义提示词", lines=15)
                
                with gr.Column(scale=1):
                    with gr.Group() as params_group:
                        gr.Markdown("### 通用参数")
                        language = gr.Dropdown(["中文", "英文"], label="语言", value="中文")
                        ratio = gr.Dropdown(["1*1", "3*2", "2*3"], label="比例", value="1*1")
                        num_images = gr.Dropdown(["1", "2", "4"], label="生图数量", value="4")
                        other_attrs = gr.Textbox(label="其他属性", value="真实风格，字体置灰")
                
                # 提示词和图像输出区
                with gr.Column(scale=1):
                    prompt_output = gr.Textbox(label="生成的提示词/输出结果", lines=15, value=prompt_default)
            image_output = gr.Gallery(label="生成的图像", columns=2, height=600)
            
            # 模板切换函数
            def show_city_scene():
                return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)
            
            def show_knowledge():
                return gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)
            
            def show_custom_template():
                return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
            
            # 生成图像函数
            def generate_images(prompt, num, ratio, curr_session_id):
                # 验证用户是否已登录
                username = app.require_auth(curr_session_id)
                # print(username)
                # print(prompt)
                logger.info(f"生图请求 - 用户: {username}, session_id: {curr_session_id}, prompt: {prompt}, num: {num}, ratio: {ratio}")
                if not username:
                    logger.warning("生图请求失败：未登录")
                    return "请先登录后再生成图像", None
            
                try:
                    # 设置图像尺寸
                    if ratio == "1*1":
                        size = "1024x1024"
                    elif ratio == "3*2":
                        size = "1792x1024"
                    elif ratio == "2*3":
                        size = "1024x1792"
                    else:
                        size = "1024x1024"
            
                    # 调用OpenAI API生成图像
                    image_urls = generate_illustration(
                        prompt=prompt,
                        n=int(num),
                        size=size,
                        model="dalle-2"
                    )
            
                    logger.info(f"生图输出 - 用户: {username}, 返回图片数量: {len(image_urls) if image_urls else 0}, urls: {image_urls}")
            
                    if not image_urls:
                        logger.error("图像生成失败，API未返回图片")
                        return "图像生成失败，请检查API设置或重试", None
            
                    return "图像生成成功", image_urls
                except Exception as e:
                    logger.exception(f"生成图像时出错: {str(e)}")
                    return f"生成图像时出错: {str(e)}", None
            
            # 按钮事件绑定
            city_scene_btn.click(
                show_city_scene,
                inputs=[],
                outputs=[city_scene_group, knowledge_group, custom_template_group]
            ).then(
                lambda: "city_scene", None, current_template  # 更新模板状态
            )

            knowledge_btn.click(
                show_knowledge,
                inputs=[],
                outputs=[city_scene_group, knowledge_group, custom_template_group]
            ).then(
                lambda: "knowledge", None, current_template
            )

            custom_template_btn.click(
                show_custom_template,
                inputs=[],
                outputs=[city_scene_group, knowledge_group, custom_template_group]
            ).then(
                lambda: "custom", None, current_template
            )

            # 生成提示词按钮事件
            def create_prompt(city, scene, topic, custom, lang, ratio, num, attrs, template_type):
                if template_type == "city_scene":
                    return PromptTemplates.city_scene_template(city, scene, lang, ratio, num, attrs)
                elif template_type == "knowledge":
                    return PromptTemplates.knowledge_template(topic, lang, ratio, num, attrs)
                elif template_type == "custom":
                    return custom
                return "请先选择一个模板"

            generate_prompt_btn.click(
                create_prompt,
                inputs=[city_input, scene_input, topic_input, custom_prompt, language, ratio, num_images, other_attrs, current_template],
                outputs=[prompt_output]
            )
            
            # 生成图像按钮事件
            generate_image_btn.click(
                generate_images,
                inputs=[prompt_output, num_images, ratio, session_id],
                outputs=[prompt_output, image_output]
            )
        
        # 原有功能选项卡保留
        with gr.Tab("功能"):
            text_input = gr.Textbox(label="输入")
            action_button = gr.Button("执行")
            result_output = gr.Textbox(label="结果")
            
            def handle_protected_function(text, curr_session_id):
                result = app.protected_function(text, curr_session_id)
                return result["message"]
            
            action_button.click(
                handle_protected_function,
                inputs=[text_input, session_id],
                outputs=[result_output]
            )

    return demo

if __name__ == "__main__":
    demo = create_gradio_interface()
    # 修改启动参数，使用默认的server_name或设置share=True
    demo.launch(share=False)